import os
import threading
import subprocess
import numpy as np
from typing import Callable, Optional
import time
import keyboard
import cv2


class RTSPFrameExtractorByFFmpeg:
    """
    RTSP流帧提取器，支持在独立线程中运行和控制
    使用FFmpeg命令行工具进行解码
    """
    def __init__(self, rtsp_url: str, frame_callback: Callable[[np.ndarray, int], None]):
        """
        初始化RTSP帧提取器

        :param rtsp_url: RTSP流地址
        :param frame_callback: 帧数据回调函数，接收参数为(frame: np.ndarray, frame_count: int)
        """
        self.rtsp_url = rtsp_url
        self.frame_callback = frame_callback
        self._thread: Optional[threading.Thread] = None
        self._running = False
        self._frame_count = 0
        self._rotate = 90
        self._ffmpeg_process: Optional[subprocess.Popen] = None

    def start(self) -> bool:
        """
        开始帧提取（在独立线程中运行）

        :return: 启动成功返回True，失败返回False
        """
        if self._running:
            print("帧提取已在运行中")
            return False

        self._running = True
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
        return True

    def stop(self) -> None:
        """
        停止帧提取并释放资源
        """
        self._running = False
        if self._thread:
            self._thread.join(timeout=5.0)
            self._thread = None
        if self._ffmpeg_process:
            self._ffmpeg_process.terminate()
            try:
                self._ffmpeg_process.wait(timeout=1)
            except subprocess.TimeoutExpired:
                self._ffmpeg_process.kill()
            self._ffmpeg_process = None
        print(f"帧提取已停止，共处理 {self._frame_count} 帧")
        self._frame_count = 0

    def _run(self) -> None:
        """
        帧提取线程主函数
        """
        # 构建FFmpeg命令
        command =  [
        'ffmpeg',
        '-hwaccel', 'cuda',              # 启用 CUDA 硬件加速
        '-c:v', 'h264_cuvid',            # 使用 NVIDIA 硬件解码器
        # '-timestamp_mode', 'preserve',  # 保留原始时间戳
        '-re',# 以本地帧率读取输入，稳定时间戳
        '-i', self.rtsp_url,                  # RTSP 输入流
        '-vf', 'format=nv12,scale=-1:1080,format=bgr24',  # 优化的滤镜链
        '-f', 'image2pipe',              # 输出为管道流
        '-pix_fmt', 'bgr24',             # 输出像素格式
        '-vcodec', 'rawvideo',           # 输出原始视频数据
        '-an', '-',                      # 忽略音频，输出到标准输出
    ]

        try:
            print(command)
            # 启动FFmpeg进程，移除 stderr=subprocess.DEVNULL 以打印 ffmpeg 日志
            self._ffmpeg_process = subprocess.Popen(
                command,
                stdout=subprocess.PIPE,
                stderr=subprocess.DEVNULL
            )
            print(f"开始从RTSP流提取帧: {self.rtsp_url}")

            while self._running:
                # 读取一帧数据
                raw_image = self._ffmpeg_process.stdout.read(1920 * 1080 * 3)  # 假设分辨率为1920x1080
                if not raw_image:
                    print("RTSP流读取失败或已结束")
                    break

                # 将原始数据转换为numpy数组
                image_array = np.frombuffer(raw_image, dtype='uint8')
                # 将一维的原始像素数据重塑为 1080x1920 的三通道图像帧
                frame = image_array.reshape((1080, 1920, 3))
                # 旋转
                frame = cv2.rotate(frame, cv2.ROTATE_180)
                # 镜像
                # frame = cv2.flip(frame, 1)


                # 调用回调函数传递当前帧和帧计数
                self.frame_callback(frame, self._frame_count)
                self._frame_count += 1

        except Exception as e:
            print(f"处理RTSP流时出错: {e}")
        finally:
            # 确保资源释放
            if self._ffmpeg_process:
                self._ffmpeg_process.terminate()
                try:
                    self._ffmpeg_process.wait(timeout=1)
                except subprocess.TimeoutExpired:
                    self._ffmpeg_process.kill()
                self._ffmpeg_process = None
            self._running = False

# 以下代码在示例使用部分添加按 esc 停止的功能
if __name__ == "__main__":
    # 用于控制程序是否继续运行的标志
    running = True
    # rtsp_url = "rtsp://admin:XG123456@192.168.1.154:554/streaming/channels/101"
    rtsp_url = "rtsp://admin:xg123456@10.10.18.168:554/streaming/channels/101"

    def handle_frame(frame: np.ndarray, frame_count: int):
        """帧处理回调函数示例，使用窗体展示图片"""
        global running
        print(f"处理第 {frame_count} 帧，尺寸: {frame.shape}")
        cv2.imshow('RTSP Frame', frame)
        # 检查是否按下 'ESC' 键（ASCII码为27）
        if cv2.waitKey(1) == 27:
            cv2.destroyAllWindows()
            running = False

    # 创建提取器实例
    extractor = RTSPFrameExtractorByFFmpeg(
        rtsp_url,
        frame_callback=handle_frame
    )

    # 启动提取
    extractor.start()

    try:
        # 等待提取器停止或用户按下ESC键
        while running and extractor._running:
            keyboard.wait('esc')
            running = False
    finally:
        # 确保停止提取器并释放资源
        extractor.stop()
        cv2.destroyAllWindows()
